

class ModelFactory():
    def __init__(self, args, cfg, **kwargs):
        self.args = args
        self.cfg = cfg

        self.class_corr = kwargs['class_corr']
    
    def build_model(self):
        model_name = self.cfg.MODEL.NAME
        model_name = model_name.upper()
        if model_name == 'UPT':
            from models.upt.model import build_detector
            return build_detector(self.cfg, self.args, class_corr=self.class_corr)
        elif model_name == 'UPT_FASTERRCNN':
            from models.upt_fasterrcnn.model import build_upt_facter_rcnn
            return build_upt_facter_rcnn(self.cfg, self.args, class_corr=self.class_corr)
        elif model_name == 'UPT_POSE':
            from models.upt_pose.model import build_detector
            return build_detector(self.cfg, self.args, class_corr=self.class_corr)
        else:
            TypeError(f'Do not supported {model_name} model now!!')
